import numpy as np

from sklearn import datasets
import random

iris = datasets.load_iris()
dataset = [(iris.data[i][None, ...], iris.target[i]) for i in range(len(iris.target))]

''' Данная нейросеть работает с batch-ами

    batch это группа образцов для обучения 
    ошибка вычисляется суммарной ошибкой для группы образцов '''

INPUT_DIM = 4  # входные параметры
OUTPUT_DIM = 3  # ответы нейронной сети

H_DIM = 10  # Количество нейронов в слое


# функции

def relu(x):  # функция активации 1 скрытого слоя
    return np.maximum(0, x)


def softmax(x):  # функция активации выходного слоя
    out = np.exp(x)
    return out / np.sum(out)


def softmax_batch(x):
    out = np.exp(x)
    return out / np.sum(out, axis=1, keepdims=True)


def sparse_cross_entropy_batch(a, b):
    return -np.log(np.array([a[j, b[j]] for j in range(len(b))]))


def to_full_batch(a, num_class):
    y_full = np.zeros((len(a), num_class))
    for j, yj in enumerate(a):
        y_full[j, yj] = 1
    return y_full


def relu_deriv(t):  # производная от функция ReLU
    return (t >= 0).astype(float)


# два полносвязных слоя
# объявление весов и bias
W1 = np.random.rand(INPUT_DIM, H_DIM)  # размерность (INPUT_DIM, H_DIM)
b1 = np.random.rand(1, H_DIM)
W2 = np.random.rand(H_DIM, OUTPUT_DIM)
b2 = np.random.rand(1, OUTPUT_DIM)

W1 = (W1 - 0.5) * 2 * np.sqrt(1 / INPUT_DIM)
b1 = (b1 - 0.5) * 2 * np.sqrt(1 / INPUT_DIM)
W2 = (W2 - 0.5) * 2 * np.sqrt(1 / H_DIM)
b2 = (b2 - 0.5) * 2 * np.sqrt(1 / H_DIM)

ALPHA = 0.001  # скорость обучения
EPOCH = 400  # количество эпох
BATCH = 50

loss_arr = []

for epoh in range(EPOCH):

    random.shuffle(dataset)  # Перемешивает dataset случайным образом.

    for i in range(len(dataset) // BATCH):  # итерации происходят не по конкретному batch, а итерируемся по batch-ам
        # * - распаковка в аргументы функции
        batch_x, batch_y = zip(*dataset[i * BATCH: i * BATCH + BATCH])
        x = np.concatenate(batch_x, axis=0)  # создаем массив из массивов batch-ов
        y = np.array(batch_y)

        # Forward

        ''' в качестве промежутчных значений мы будем получать не векторы, а матрицы
        т.е по сути параллельное вычисление каждого вектора из batch независимо '''

        t1 = x @ W1 + b1  # вычисление 1 скрытого слоя
        h1 = relu(t1)  # функция активации промежу точного выходного слоя
        t2 = h1 @ W2 + b2  # вычисление 2 скрытого слоя
        z = softmax_batch(t2)  # функция активации выходного слоя

        E = np.sum(sparse_cross_entropy_batch(z, y))  # ошибка

        # Backward

        y_full = to_full_batch(y, OUTPUT_DIM)  # правильный ответ записываем в ветор той же размерности что и output_dim

        ''' вычисление градиентов 
        dE/dt1 <- dE/dh1 <- dE/dt2 <- E
        из dE/dt2 можем получить dE/dW2 и dE/db2
        из dE/dt1 можем получить dE/dW1 и dE/db1 '''

        dE_dt2 = z - y_full
        dE_dW2 = h1.T @ dE_dt2
        dE_b2 = np.sum(dE_dt2, axis=0, keepdims=True)

        dE_h1 = dE_dt2 @ W2.T

        dE_dt1 = dE_h1 * relu_deriv(t1)
        dE_dW1 = x.T @ dE_dt1
        dE_db1 = np.sum(dE_dt1, axis=0, keepdims=True)

        # Update

        W1 -= ALPHA * dE_dW1
        b1 -= ALPHA * dE_db1
        W2 -= ALPHA * dE_dW2
        b2 -= ALPHA * dE_b2

        loss_arr.append(E)


def predict(x):
    t1 = x @ W1 + b1  # @ - матричное умножение
    h1 = relu(t1)
    t2 = h1 @ W2 + b2
    z = softmax(t2)
    return z


def calc_accuracy():
    correct = 0
    for x, y in dataset:
        z = predict(x)
        y_pred = np.argmax(z)
        if y_pred == y:
            correct += 1
    acc = correct / len(dataset)
    return acc


accuracy = calc_accuracy()
print("Accuracy: ", accuracy)

import matplotlib.pyplot as plt

plt.plot(loss_arr)
plt.show()
